from train_api import *
from utils.parsers import parse_cfg
import argparse
import torch


def main():
    parser = argparse.ArgumentParser(description="PyTorch Object Detection Training")
    parser.add_argument(
        "--configfile",
        help="path to config file",
        default='F:/detect_change/cfg/unet/train_unet.cfg',
        type=str,
    )

    args = parser.parse_args()
    cfgs = parse_cfg(args.configfile)

    torch.cuda.set_device(cfgs["gpus"])
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if cfgs["model_type"] == "unet":
        train_unet(cfgs,device)


if __name__ == '__main__':
    main()
